package com.zhentao.ai.algorithm;

import com.zhentao.ai.model.ConfidenceInterval;
import com.zhentao.ai.model.DataPoint;
import com.zhentao.ai.model.PredictionResult;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.List;

/**
 * ARIMA时间序列预测算法实现
 * Auto-Regressive Integrated Moving Average
 * 
 * @author zhentao
 * @date 2024-12-19
 */
@Slf4j
@Component
public class ARIMAPredictor {
    
    private static final int DEFAULT_P = 1; // AR项数
    private static final int DEFAULT_D = 1; // 差分次数
    private static final int DEFAULT_Q = 1; // MA项数
    
    /**
     * ARIMA预测主方法
     * @param historicalData 历史数据点
     * @param forecastPeriods 预测期数
     * @param confidenceLevel 置信水平 (90, 95, 99)
     * @return 预测结果
     */
    public PredictionResult predict(List<DataPoint> historicalData, int forecastPeriods, int confidenceLevel) {
        if (historicalData == null || historicalData.size() < 10) {
            throw new IllegalArgumentException("历史数据不足，至少需要10个数据点");
        }
        
        log.info("开始ARIMA预测，历史数据点数: {}, 预测期数: {}", historicalData.size(), forecastPeriods);
        
        // 1. 数据预处理
        double[] values = preprocessData(historicalData);
        
        // 2. 检验平稳性并差分
        double[] stationaryData = makeStationary(values);
        
        // 3. 参数估计 (简化版本)
        ARIMAModel model = estimateParameters(stationaryData, DEFAULT_P, DEFAULT_D, DEFAULT_Q);
        
        // 4. 进行预测
        double[] predictions = forecast(model, stationaryData, forecastPeriods);
        
        // 5. 反差分还原
        double[] finalPredictions = reverseStationarity(predictions, values);
        
        // 6. 计算置信区间
        ConfidenceInterval[] intervals = calculateConfidenceIntervals(finalPredictions, model.residualVariance, confidenceLevel);
        
        // 7. 构建预测结果
        return buildPredictionResult(finalPredictions, intervals, model);
    }
    
    /**
     * 数据预处理 - 提取数值并处理异常值
     */
    private double[] preprocessData(List<DataPoint> data) {
        double[] values = new double[data.size()];
        for (int i = 0; i < data.size(); i++) {
            values[i] = data.get(i).getValue().doubleValue();
        }
        
        // 处理异常值 - 使用中位数替换极值
        double median = calculateMedian(values);
        double mad = calculateMAD(values, median); // 中位数绝对偏差
        double threshold = 3.0 * mad;
        
        for (int i = 0; i < values.length; i++) {
            if (Math.abs(values[i] - median) > threshold) {
                values[i] = median;
                log.warn("检测到异常值，已用中位数替换: 位置{}, 原值{}, 新值{}", i, data.get(i).getValue(), median);
            }
        }
        
        return values;
    }
    
    /**
     * 平稳性检验和差分处理
     */
    private double[] makeStationary(double[] data) {
        // 简化版ADF检验 - 实际项目中应使用更严格的统计检验
        if (isStationary(data)) {
            log.info("数据已平稳，无需差分");
            return data;
        }
        
        // 一阶差分
        double[] diff = new double[data.length - 1];
        for (int i = 1; i < data.length; i++) {
            diff[i - 1] = data[i] - data[i - 1];
        }
        
        log.info("执行一阶差分，数据长度: {} -> {}", data.length, diff.length);
        return diff;
    }
    
    /**
     * 简化的平稳性检验
     */
    private boolean isStationary(double[] data) {
        if (data.length < 20) return false;
        
        // 计算滑动平均的方差
        int windowSize = Math.min(data.length / 4, 20);
        double[] movingMeans = new double[data.length - windowSize + 1];
        
        for (int i = 0; i <= data.length - windowSize; i++) {
            double sum = 0;
            for (int j = i; j < i + windowSize; j++) {
                sum += data[j];
            }
            movingMeans[i] = sum / windowSize;
        }
        
        double variance = calculateVariance(movingMeans);
        return variance < 0.1; // 简化判断标准
    }
    
    /**
     * ARIMA参数估计
     */
    private ARIMAModel estimateParameters(double[] data, int p, int d, int q) {
        ARIMAModel model = new ARIMAModel(p, d, q);
        
        // 简化的最小二乘估计 - 实际应使用MLE或条件最小二乘
        if (data.length > p + q) {
            // AR系数估计 (简化版本)
            model.arCoefficients = estimateARCoefficients(data, p);
            
            // MA系数估计 (简化版本)
            model.maCoefficients = estimateMACoefficients(data, model.arCoefficients, q);
            
            // 残差方差估计
            model.residualVariance = calculateResidualVariance(data, model);
        } else {
            // 数据不足时使用默认参数
            model.arCoefficients = new double[]{0.5};
            model.maCoefficients = new double[]{0.3};
            model.residualVariance = calculateVariance(data);
        }
        
        log.info("ARIMA({},{},{}) 参数估计完成, 残差方差: {}", p, d, q, model.residualVariance);
        return model;
    }
    
    /**
     * AR系数估计 - 使用Yule-Walker方程
     */
    private double[] estimateARCoefficients(double[] data, int p) {
        if (p == 0) return new double[0];
        
        double[] coeffs = new double[p];
        double mean = calculateMean(data);
        
        // 计算自相关函数
        double[] autocorr = new double[p + 1];
        for (int lag = 0; lag <= p; lag++) {
            autocorr[lag] = calculateAutocorrelation(data, lag, mean);
        }
        
        // 解Yule-Walker方程 (简化版本)
        if (p == 1) {
            coeffs[0] = autocorr[1] / autocorr[0];
        } else {
            // 多项AR系数估计 (简化处理)
            for (int i = 0; i < p; i++) {
                coeffs[i] = autocorr[i + 1] / (autocorr[0] + 0.001); // 避免除零
            }
        }
        
        return coeffs;
    }
    
    /**
     * MA系数估计
     */
    private double[] estimateMACoefficients(double[] data, double[] arCoeffs, int q) {
        if (q == 0) return new double[0];
        
        // 简化的MA系数估计
        double[] maCoeffs = new double[q];
        for (int i = 0; i < q; i++) {
            maCoeffs[i] = 0.1 + 0.1 * i; // 简化初始值
        }
        
        return maCoeffs;
    }
    
    /**
     * 预测核心算法
     */
    private double[] forecast(ARIMAModel model, double[] data, int periods) {
        double[] predictions = new double[periods];
        double[] extendedData = new double[data.length + periods];
        
        // 复制历史数据
        System.arraycopy(data, 0, extendedData, 0, data.length);
        
        for (int i = 0; i < periods; i++) {
            int currentIndex = data.length + i;
            double prediction = 0.0;
            
            // AR部分
            for (int j = 0; j < model.arCoefficients.length; j++) {
                if (currentIndex - j - 1 >= 0) {
                    prediction += model.arCoefficients[j] * extendedData[currentIndex - j - 1];
                }
            }
            
            // MA部分 (简化处理，实际应考虑残差)
            for (int j = 0; j < model.maCoefficients.length; j++) {
                prediction += model.maCoefficients[j] * Math.sqrt(model.residualVariance) * 0.1;
            }
            
            extendedData[currentIndex] = prediction;
            predictions[i] = prediction;
        }
        
        return predictions;
    }
    
    /**
     * 反差分还原
     */
    private double[] reverseStationarity(double[] diffPredictions, double[] originalData) {
        double[] restored = new double[diffPredictions.length];
        double lastValue = originalData[originalData.length - 1];
        
        for (int i = 0; i < diffPredictions.length; i++) {
            restored[i] = lastValue + diffPredictions[i];
            lastValue = restored[i]; // 累积效应
        }
        
        return restored;
    }
    
    /**
     * 计算置信区间
     */
    private ConfidenceInterval[] calculateConfidenceIntervals(double[] predictions, double variance, int confidenceLevel) {
        ConfidenceInterval[] intervals = new ConfidenceInterval[predictions.length];
        
        // 根据置信水平确定z值
        double zValue;
        switch (confidenceLevel) {
            case 90: zValue = 1.645; break;
            case 95: zValue = 1.96; break;
            case 99: zValue = 2.576; break;
            default: zValue = 1.96;
        }
        
        double stdError = Math.sqrt(variance);
        
        for (int i = 0; i < predictions.length; i++) {
            double margin = zValue * stdError * Math.sqrt(i + 1); // 预测误差随时间增长
            intervals[i] = new ConfidenceInterval(
                BigDecimal.valueOf(predictions[i] - margin).setScale(2, RoundingMode.HALF_UP),
                BigDecimal.valueOf(predictions[i] + margin).setScale(2, RoundingMode.HALF_UP),
                confidenceLevel
            );
        }
        
        return intervals;
    }
    
    // 辅助数学方法
    private double calculateMean(double[] data) {
        double sum = 0;
        for (double value : data) sum += value;
        return sum / data.length;
    }
    
    private double calculateVariance(double[] data) {
        double mean = calculateMean(data);
        double sum = 0;
        for (double value : data) {
            sum += Math.pow(value - mean, 2);
        }
        return sum / (data.length - 1);
    }
    
    private double calculateMedian(double[] data) {
        double[] sorted = data.clone();
        java.util.Arrays.sort(sorted);
        int n = sorted.length;
        return n % 2 == 0 ? (sorted[n/2-1] + sorted[n/2]) / 2.0 : sorted[n/2];
    }
    
    private double calculateMAD(double[] data, double median) {
        double[] deviations = new double[data.length];
        for (int i = 0; i < data.length; i++) {
            deviations[i] = Math.abs(data[i] - median);
        }
        return calculateMedian(deviations);
    }
    
    private double calculateAutocorrelation(double[] data, int lag, double mean) {
        double numerator = 0;
        double denominator = 0;
        
        for (int i = lag; i < data.length; i++) {
            numerator += (data[i] - mean) * (data[i - lag] - mean);
        }
        
        for (int i = 0; i < data.length; i++) {
            denominator += Math.pow(data[i] - mean, 2);
        }
        
        return denominator != 0 ? numerator / denominator : 0;
    }
    
    private double calculateResidualVariance(double[] data, ARIMAModel model) {
        // 简化的残差方差计算
        return calculateVariance(data) * 0.8; // 假设模型解释了20%的方差
    }
    
    private PredictionResult buildPredictionResult(double[] predictions, ConfidenceInterval[] intervals, ARIMAModel model) {
        PredictionResult result = new PredictionResult();
        result.setPredictions(predictions);
        result.setConfidenceIntervals(intervals);
        result.setModelInfo(String.format("ARIMA(%d,%d,%d)", model.p, model.d, model.q));
        result.setResidualVariance(BigDecimal.valueOf(model.residualVariance));
        return result;
    }
    
    /**
     * ARIMA模型内部类
     */
    private static class ARIMAModel {
        int p, d, q;
        double[] arCoefficients;
        double[] maCoefficients;
        double residualVariance;
        
        public ARIMAModel(int p, int d, int q) {
            this.p = p;
            this.d = d;
            this.q = q;
        }
    }
}
